#!/bin/bash
set -e

# Directory setup
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
cd "$SCRIPT_DIR"

# Create output directories if they don't exist
mkdir -p pngs

# Check if compiled binary exists
if [ ! -f "./rsdc" ]; then
    echo "Compiled binary not found. Running compilation script..."
    ./rscompile.sh
fi

# Run the capture program
echo "Starting arm training capture process..."
echo "You will be prompted to enter target arm positions."
echo "After the arm motion completes, press Ctrl+C to stop capture and save data."
./rsdc

# Process the captured data if motion_data.json exists
if [ -f "motion_data.json" ]; then
    echo "Motion data found. Processing captured frames..."
    
    # Create a directory for this capture session with timestamp
    SESSION_DIR="session_$(date +%Y%m%d_%H%M%S)"
    mkdir -p "$SESSION_DIR"
    
    # Copy motion data to session directory
    cp motion_data.json "$SESSION_DIR/"
    
    # Move color frames to session directory
    if [ -d "pngs" ] && [ "$(ls -A pngs)" ]; then
        mkdir -p "$SESSION_DIR/color_frames"
        cp pngs/* "$SESSION_DIR/color_frames/"
        echo "Copied $(ls pngs | wc -l) color frames to $SESSION_DIR/color_frames/"
    fi
    
    # Create a summary file
    echo "Arm motion summary:" > "$SESSION_DIR/summary.txt"
    echo "Session: $SESSION_DIR" >> "$SESSION_DIR/summary.txt"
    echo "Time: $(date)" >> "$SESSION_DIR/summary.txt"
    echo "Number of frames: $(ls pngs | wc -l)" >> "$SESSION_DIR/summary.txt"
    echo "See motion_data.json for detailed joint movement data" >> "$SESSION_DIR/summary.txt"
    
    echo "Training session data saved to $SESSION_DIR/"
    echo "You can review the motion data in $SESSION_DIR/motion_data.json"
else
    echo "No motion data was generated. The capture may have been interrupted before arm movement."
fi

echo "Arm training process complete." 